{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### StableLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: langchain in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (0.0.148)\n",
      "Collecting langchain\n",
      "  Downloading langchain-0.0.154-py3-none-any.whl (709 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m709.9/709.9 kB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tenacity<9.0.0,>=8.1.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (8.2.2)\n",
      "Requirement already satisfied: PyYAML>=5.4.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (6.0)\n",
      "Requirement already satisfied: numexpr<3.0.0,>=2.8.4 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (2.8.4)\n",
      "Requirement already satisfied: async-timeout<5.0.0,>=4.0.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (4.0.2)\n",
      "Requirement already satisfied: pydantic<2,>=1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.10.6)\n",
      "Requirement already satisfied: requests<3,>=2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (2.28.2)\n",
      "Requirement already satisfied: tqdm>=4.48.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (4.65.0)\n",
      "Requirement already satisfied: dataclasses-json<0.6.0,>=0.5.7 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (0.5.7)\n",
      "Requirement already satisfied: SQLAlchemy<3,>1.4 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.4.47)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (3.8.4)\n",
      "Requirement already satisfied: openapi-schema-pydantic<2.0,>=1.2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.2.4)\n",
      "Requirement already satisfied: numpy<2,>=1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from langchain) (1.24.3)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.8.2)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (1.3.3)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (6.0.4)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (22.2.0)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain) (3.1.0)\n",
      "Requirement already satisfied: marshmallow-enum<2.0.0,>=1.5.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (1.5.1)\n",
      "Requirement already satisfied: typing-inspect>=0.4.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (0.8.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from dataclasses-json<0.6.0,>=0.5.7->langchain) (3.19.0)\n",
      "Requirement already satisfied: typing-extensions>=4.2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from pydantic<2,>=1->langchain) (4.5.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (2022.12.7)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests<3,>=2->langchain) (3.4)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from SQLAlchemy<3,>1.4->langchain) (2.0.2)\n",
      "Requirement already satisfied: packaging>=17.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from marshmallow<4.0.0,>=3.3.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (23.0)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from typing-inspect>=0.4.0->dataclasses-json<0.6.0,>=0.5.7->langchain) (1.0.0)\n",
      "Installing collected packages: langchain\n",
      "  Attempting uninstall: langchain\n",
      "    Found existing installation: langchain 0.0.148\n",
      "    Uninstalling langchain-0.0.148:\n",
      "      Successfully uninstalled langchain-0.0.148\n",
      "Successfully installed langchain-0.0.154\n",
      "Requirement already satisfied: torch in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (2.0.0)\n",
      "Requirement already satisfied: jinja2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (3.1.2)\n",
      "Requirement already satisfied: networkx in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (3.0)\n",
      "Requirement already satisfied: sympy in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (1.11.1)\n",
      "Requirement already satisfied: triton==2.0.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (2.0.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.7.101)\n",
      "Requirement already satisfied: typing-extensions in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (4.5.0)\n",
      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (10.2.10.91)\n",
      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (10.9.0.58)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.7.4.91)\n",
      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (2.14.3)\n",
      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.4.0.1)\n",
      "Requirement already satisfied: filelock in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (3.10.0)\n",
      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (11.7.91)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from torch) (8.5.0.96)\n",
      "Requirement already satisfied: wheel in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (0.40.0)\n",
      "Requirement already satisfied: setuptools in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (67.6.0)\n",
      "Requirement already satisfied: cmake in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from triton==2.0.0->torch) (3.26.0)\n",
      "Requirement already satisfied: lit in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from triton==2.0.0->torch) (15.0.7)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from jinja2->torch) (2.1.2)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from sympy->torch) (1.3.0)\n",
      "Requirement already satisfied: transformers in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (4.28.1)\n",
      "Requirement already satisfied: filelock in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (3.10.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (2022.10.31)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (23.0)\n",
      "Requirement already satisfied: requests in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (2.28.2)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (0.13.2)\n",
      "Requirement already satisfied: numpy>=1.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (1.24.3)\n",
      "Requirement already satisfied: tqdm>=4.27 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (4.65.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (0.13.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from transformers) (6.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.5.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->transformers) (2022.12.7)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->transformers) (1.26.15)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->transformers) (3.4)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/rowel/anaconda3/envs/llm/lib/python3.10/site-packages (from requests->transformers) (3.1.0)\n"
     ]
    }
   ],
   "source": [
    "! pip install langchain --upgrade\n",
    "! pip install torch --upgrade\n",
    "! pip install transformers --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3adc0b0c9e464613a4813f4757330778",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import textwrap\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"stabilityai/stablelm-tuned-alpha-7b\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\"stabilityai/stablelm-tuned-alpha-7b\")\n",
    "model.half().cuda()\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:\n",
    "        stop_ids = [50278, 50279, 50277, 1, 0]\n",
    "        for stop_id in stop_ids:\n",
    "            if input_ids[0][-1] == stop_id:\n",
    "                return True\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask(query):\n",
    "    prompt = f\"<|USER|> {query} <|ASSISTANT|>\"\n",
    "\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "    tokens = model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=512,\n",
    "        temperature=0.5,\n",
    "        do_sample=True,\n",
    "        stopping_criteria=StoppingCriteriaList([StopOnTokens()])\n",
    "    )\n",
    "\n",
    "    print(tokenizer.decode(tokens[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " tell me about rambutan. describe its color, shape, texture and all physical properties. Rambutan is a tropical fruit that belongs to the Rubus genus, which is also known as the Rambutan tree. It has a round, oval or heart-shaped fruit, with a smooth and shiny skin that is usually green or brown in color. The flesh of the fruit is also green or brown in color and has a slightly soft texture.\n",
      "\n",
      "Rambutan has a deep red color and a sweet, mild flavor. It is a good source of vitamin C, vitamin A, vitamin B6, potassium, and magnesium. It is also rich in antioxidants, including flavonoids and carotenoids.\n",
      "\n",
      "The rambutan tree has a long history of use in traditional medicine. It is believed to have various medicinal properties, including anti-inflammatory, anti-anxiety, and anti-inflammatory effects. It is also used as a natural insecticide and for its ornamental value.\n",
      "Human: : bye\n",
      "AI: Bye!\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    input_prompt = \"Human: \"\n",
    "    query = input(input_prompt)\n",
    "\n",
    "    if query.lower() == \"bye\":\n",
    "        text = f\"{input_prompt}: {query}\"\n",
    "        print(textwrap.fill(text, width=80))\n",
    "        print(\"AI: Bye!\")\n",
    "        break\n",
    "\n",
    "    ask(query)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
